#!/usr/bin/env python3

import argparse
from glob import glob
import os
import pendulum
import f90nml
from netCDF4 import Dataset
import re
from shutil import copy
import sys
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../utils')
from utils import cli, parse_config, wrf_version, has_key, Version

def get_num_land_cat(wrfinput):
	wrfinput = Dataset(wrfinput)
	num_land_cat = wrfinput.getncattr('NUM_LAND_CAT')
	wrfinput.close()
	return num_land_cat

def config_wrf(work_root, wrf_root, config, args, wps_work_dir=None, tag=None):
	phys_config = config['physics'] if 'physics' in config else {}

	start_time = config['custom']['start_time']
	end_time = config['custom']['end_time']
	max_dom = config['domains']['max_dom']
	
	start_time_str = start_time.format('YYYY-MM-DD_HH:mm:ss')
	end_time_str = end_time.format('YYYY-MM-DD_HH:mm:ss')

	if tag != None:
		wrf_work_dir = f'{work_root}/{tag}/wrf'
	else:
		wrf_work_dir = f'{work_root}/wrf'
	if not os.path.isdir(wrf_work_dir): os.makedirs(wrf_work_dir)
	old_dir = os.getcwd()
	os.chdir(wrf_work_dir)

	if wps_work_dir == None:
		wps_work_dir = f'{work_root}/wps'
	if not os.path.isdir(wps_work_dir): cli.error(f'WPS work directory {wps_work_dir} does not exist!')

	version = wrf_version(wrf_root)

	if has_key(config, ('physics', 'num_land_cat')):
		num_land_cat = config['physics']['num_land_cat']
	elif os.path.isfile(f'{wps_work_dir}/met_em.d01.{start_time_str}.nc'):
		num_land_cat = get_num_land_cat(f'{wps_work_dir}/met_em.d01.{start_time_str}.nc')
	elif os.path.isfile(f'{wrf_work_dir}/wrfinput_d01_{start_time_str}'):
		num_land_cat = get_num_land_cat(f'{wrf_work_dir}/wrfinput_d01_{start_time_str}')
	elif os.path.isfile(f'{wrf_work_dir}/wrfinput_d01'):
		num_land_cat = get_num_land_cat(f'{wrf_work_dir}/wrfinput_d01')
	elif os.path.isfile(f'{wrf_work_dir}/wrfout_d01_{start_time_str}'):
		num_land_cat = get_num_land_cat(f'{wrf_work_dir}/wrfout_d01_{start_time_str}')
	else:
		cli.warning(f'Cannot get num_land_cat parameter from an existing wrfinput_d01 file in {wrf_work_dir}!')
		num_land_cat = None

	# Setup restart run.
	if has_key(config, ('time_control', 'restart')):
		if config['time_control']['restart'] == True:
			try:
				wrfrst_file = sorted(glob('wrfrst_d02_*'))[-1]
				restart_time = pendulum.from_format(wrfrst_file[11:], 'YYYY-MM-DD_HH:mm:SS')
				cli.notice(f'Restart from {restart_time}.')
				start_time = restart_time
				config['custom']['forecast_hours'] = (end_time - restart_time).in_hours()
				if has_key(config, ('custom', 'restart_time_step')):
					config['domains']['time_step'] = config['custom']['restart_time_step']
			except Exception as e:
				config['time_control']['restart'] = False
				cli.warning(f'Set restart to true, but there is no wrfrst files! Set restart to false.\n{e}')

	cli.notice('Edit namelist.input for WRF.')
	if os.path.isfile(f'{wrf_root}/run/namelist.input'):
		copy(f'{wrf_root}/run/namelist.input', 'namelist.input')
	elif os.path.isfile(f'{wrf_root}/test/em_real/namelist.input'):
		copy(f'{wrf_root}/test/em_real/namelist.input', 'namelist.input')
	namelist_input = f90nml.read('namelist.input')
	namelist_input['time_control']['run_hours']              = config['custom']['forecast_hours']
	namelist_input['time_control']['start_year']             = [int(start_time.format("Y")) for i in range(max_dom)]
	namelist_input['time_control']['start_month']            = [int(start_time.format("M")) for i in range(max_dom)]
	namelist_input['time_control']['start_day']              = [int(start_time.format("D")) for i in range(max_dom)]
	namelist_input['time_control']['start_hour']             = [int(start_time.format("H")) for i in range(max_dom)]
	namelist_input['time_control']['end_year']               = [int(end_time.format("Y")) for i in range(max_dom)]
	namelist_input['time_control']['end_month']              = [int(end_time.format("M")) for i in range(max_dom)]
	namelist_input['time_control']['end_day']                = [int(end_time.format("D")) for i in range(max_dom)]
	namelist_input['time_control']['end_hour']               = [int(end_time.format("H")) for i in range(max_dom)]
	namelist_input['time_control']['frames_per_outfile']     = [1 for i in range(max_dom)]
	namelist_input['time_control']['input_from_file']        = [True for i in range(max_dom)]
	if 'background' in config['custom'] and 'interval_seconds' in config['custom']['background']:
		namelist_input['time_control']['interval_seconds']       = config['custom']['background']['interval_seconds']
	if 'time_control' in config:
		for key, value in config['time_control'].items():
			namelist_input['time_control'][key] = value
	for key, value in config['domains'].items():
		namelist_input['domains'][key] = value
	if 'physics_suite' in namelist_input['physics']: del namelist_input['physics']['physics_suite']
	for key, value in phys_config.items():
		namelist_input['physics'][key] = value
	if num_land_cat != None: namelist_input['physics']['num_land_cat'] = num_land_cat
	if 'dynamics' in config:
		for key, value in config['dynamics'].items():
			namelist_input['dynamics'][key] = value
	if 'fdda' in config:
		for key, value in config['fdda'].items():
			namelist_input['fdda'][key] = value
	if 'namelist_quilt' in config:
		for key, value in config['namelist_quilt'].items():
			namelist_input['namelist_quilt'][key] = value
	if 'diags' in config:
		for key, value in config['diags'].items():
			namelist_input['diags'][key] = value
	if version == Version('3.9.1'):
		namelist_input['dynamics']['gwd_opt'] = 0
	namelist_input.write('./namelist.input', force=True)

	cli.notice('Succeeded.')
	os.chdir(old_dir)
